# ingest_q2sql.py
import logging
from pymilvus import MilvusClient, DataType, FieldSchema, CollectionSchema
from pymilvus import connections, utility, Collection
import torch
import json
from sentence_transformers import SentenceTransformer

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 1. 初始化嵌入函数 - 使用本地BGE-M3模型
bge_m3_model = "D:\\code\\aicode\\sentence-transformers\\bge-m3"
embedding_model = SentenceTransformer(bge_m3_model)

def local_embedding_function(texts):
    """本地嵌入函数包装器"""
    if isinstance(texts, str):
        texts = [texts]
    embeddings = embedding_model.encode(texts, 
                                      batch_size=12, 
                                      convert_to_tensor=False,
                                      normalize_embeddings=True)
    return embeddings.tolist()

# 2. 加载 Q->SQL 对（假设 q2sql_pairs.json 数组，每项 { "question": ..., "sql": ... }）
with open("sakila/q2sql_pairs.json", "r") as f:
    pairs = json.load(f)
    logging.info(f"[Q2SQL] 从JSON文件加载了 {len(pairs)} 个问答对")

# 3. 连接 Milvus - 使用Docker部署的Milvus服务
connections.connect(
    alias="default",
    host='localhost',     # Milvus服务器地址
    port='19530'          # Milvus服务器端口
)
# 检查连接状态
logging.info(f"[Q2SQL] 成功连接到Milvus服务器: {utility.list_collections()}")


# 4. 定义 Collection Schema
# 获取向量维度，添加错误处理
try:
    dummy_embedding = local_embedding_function(["dummy test"])[0]
    vector_dim = len(dummy_embedding)
    logging.info(f"[Q2SQL] 检测到向量维度: {vector_dim}")
except Exception as e:
    logging.error(f"[Q2SQL] 获取向量维度失败: {e}")
    # BGE-M3 默认维度是 1024
    vector_dim = 1024
    logging.info(f"[Q2SQL] 使用默认向量维度: {vector_dim}")

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
    FieldSchema(name="question", dtype=DataType.VARCHAR, max_length=500),
    FieldSchema(name="sql_text", dtype=DataType.VARCHAR, max_length=2000),
]
schema = CollectionSchema(fields, description="Q2SQL Knowledge Base", enable_dynamic_field=False)

# 5. 创建 Collection（如不存在）
collection_name = "q2sql_knowledge"

if utility.has_collection(collection_name):
    logging.info(f"[Q2SQL] 集合 {collection_name} 已存在，正在删除...")
    utility.drop_collection(collection_name)
    logging.info(f"[Q2SQL] 成功删除集合 {collection_name}")    

if not utility.has_collection(collection_name):
    collection = Collection(name=collection_name, schema=schema)
    logging.info(f"[Q2SQL] 创建了新的集合 {collection_name}")
else:
    logging.info(f"[Q2SQL] 集合 {collection_name} 已存在")

# 6. 为向量字段添加索引
index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "COSINE",
    "params": {"nlist": 1024}
}

collection.create_index(field_name="vector", index_params=index_params)

# 7. 批量插入 Q2SQL 对
data = []
texts = []
for pair in pairs:
    texts.append(pair["question"])
    data.append({"question": pair["question"], "sql_text": pair["sql"]})

logging.info(f"[Q2SQL] 准备处理 {len(data)} 个问答对")

# 生成全部嵌入
embeddings = local_embedding_function(texts)
logging.info(f"[Q2SQL] 成功生成了 {len(embeddings)} 个向量嵌入")

# 组织为 Milvus insert 格式
records = []
for emb, rec in zip(embeddings, data):
    rec["vector"] = emb
    records.append(rec)

res = collection.insert(data=records)
logging.info(f"[Q2SQL] 成功插入了 {len(records)} 条记录到Milvus")
logging.info(f"[Q2SQL] 插入结果: {res}")

# 确保数据持久化
collection.flush()
logging.info("[Q2SQL] 知识库构建完成")
